package main

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/xuri/excelize/v2"
)

// ETFData 存储ETF周线数据
type ETFData struct {
	Date    string
	Open    float64
	Close   float64
	High    float64
	Low     float64
	Volume  float64
	Amount  float64
	ETFCode string
}

// EastMoneyResponse 东方财富API响应结构
type EastMoneyResponse struct {
	Data struct {
		Code   string   `json:"code"`
		Klines []string `json:"klines"`
	} `json:"data"`
}

// ETFScraper ETF数据爬虫
type ETFScraper struct {
	client *http.Client
}

// NewETFScraper 创建爬虫实例
func NewETFScraper() *ETFScraper {
	return &ETFScraper{
		client: &http.Client{
			Timeout: 30 * time.Second,
		},
	}
}

// GetETFData 获取ETF数据（支持周线和日线）
func (s *ETFScraper) GetETFData(etfCode, startDate, endDate string, isDaily bool) ([]ETFData, error) {
	// 判断市场代码 (1=沪市, 0=深市)
	marketCode := "0"
	if strings.HasPrefix(etfCode, "5") {
		marketCode = "1"
	}

	// klt: 101=日线, 102=周线
	klt := "102"
	if isDaily {
		klt = "101"
	}

	// 构建URL
	url := fmt.Sprintf(
		"http://push2his.eastmoney.com/api/qt/stock/kline/get?secid=%s.%s&fields1=f1,f2,f3,f4,f5,f6&fields2=f51,f52,f53,f54,f55,f56,f57,f58,f59,f60,f61&klt=%s&fqt=1&beg=%s&end=%s&ut=fa5fd1943c7b386f172d6893dbfba10b&lmt=10000",
		marketCode, etfCode, klt, startDate, endDate,
	)

	// 发送请求
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Set("User-Agent", "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36")
	req.Header.Set("Referer", "http://quote.eastmoney.com/")

	resp, err := s.client.Do(req)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	// 解析JSON
	var result EastMoneyResponse
	if err := json.Unmarshal(body, &result); err != nil {
		return nil, err
	}

	if len(result.Data.Klines) == 0 {
		return nil, fmt.Errorf("ETF %s 没有返回数据", etfCode)
	}

	// 解析K线数据
	var data []ETFData
	for _, line := range result.Data.Klines {
		parts := strings.Split(line, ",")
		if len(parts) < 7 {
			continue
		}

		open, _ := strconv.ParseFloat(parts[1], 64)
		close, _ := strconv.ParseFloat(parts[2], 64)
		high, _ := strconv.ParseFloat(parts[3], 64)
		low, _ := strconv.ParseFloat(parts[4], 64)
		volume, _ := strconv.ParseFloat(parts[5], 64)
		amount, _ := strconv.ParseFloat(parts[6], 64)

		data = append(data, ETFData{
			Date:    parts[0],
			Open:    open,
			Close:   close,
			High:    high,
			Low:     low,
			Volume:  volume,
			Amount:  amount,
			ETFCode: etfCode,
		})
	}

	dataType := "周线"
	if isDaily {
		dataType = "日线"
	}
	fmt.Printf("成功获取 ETF %s 的 %d 条%s数据\n", etfCode, len(data), dataType)
	return data, nil
}

// GetETFWeeklyData 获取ETF周线数据（兼容老接口）
func (s *ETFScraper) GetETFWeeklyData(etfCode, startDate, endDate string) ([]ETFData, error) {
	return s.GetETFData(etfCode, startDate, endDate, false)
}

// GetETFDailyData 获取ETF日线数据
func (s *ETFScraper) GetETFDailyData(etfCode, startDate, endDate string) ([]ETFData, error) {
	return s.GetETFData(etfCode, startDate, endDate, true)
}

// ScrapeAllETFs 爬取所有ETF数据并保存
func (s *ETFScraper) ScrapeAllETFs(etfCodes []string, startDate, endDate, outputFile string, isDaily bool) error {
	var allData []ETFData

	var wg sync.WaitGroup
	var mu sync.Mutex
	for _, code := range etfCodes {
		code := code
		wg.Add(1)
		go func() {
			defer wg.Done()
			fmt.Printf("\n正在获取 ETF %s 的数据...\n", code)
			data, err := s.GetETFData(code, startDate, endDate, isDaily)
			if err != nil {
				fmt.Printf("获取 ETF %s 数据失败: %v\n", code, err)
				return
			}

			mu.Lock()
			defer mu.Unlock()
			allData = append(allData, data...)
		}()
	}

	wg.Wait()

	if len(allData) == 0 {
		return fmt.Errorf("没有获取到任何数据")
	}

	// 保存到Excel
	if err := s.SaveToExcel(allData, outputFile, isDaily); err != nil {
		return err
	}

	fmt.Printf("\n数据已保存到 %s\n", outputFile)
	fmt.Printf("总共 %d 条记录\n", len(allData))

	return nil
}

// SaveToExcel 保存数据到Excel
func (s *ETFScraper) SaveToExcel(data []ETFData, filename string, isDaily bool) error {
	f := excelize.NewFile()
	defer f.Close()

	sheetName := "ETF周线数据"
	if isDaily {
		sheetName = "ETF日线数据"
	}
	index, err := f.NewSheet(sheetName)
	if err != nil {
		return err
	}

	// 设置表头
	headers := []string{"日期", "开盘价", "收盘价", "最高价", "最低价", "成交量", "成交额", "ETF代码"}
	for i, header := range headers {
		cell := fmt.Sprintf("%s1", string(rune('A'+i)))
		f.SetCellValue(sheetName, cell, header)
	}

	// 写入数据
	for i, record := range data {
		row := i + 2
		f.SetCellValue(sheetName, fmt.Sprintf("A%d", row), record.Date)
		f.SetCellValue(sheetName, fmt.Sprintf("B%d", row), record.Open)
		f.SetCellValue(sheetName, fmt.Sprintf("C%d", row), record.Close)
		f.SetCellValue(sheetName, fmt.Sprintf("D%d", row), record.High)
		f.SetCellValue(sheetName, fmt.Sprintf("E%d", row), record.Low)
		f.SetCellValue(sheetName, fmt.Sprintf("F%d", row), record.Volume)
		f.SetCellValue(sheetName, fmt.Sprintf("G%d", row), record.Amount)
		f.SetCellValue(sheetName, fmt.Sprintf("H%d", row), record.ETFCode)
	}

	f.SetActiveSheet(index)
	f.DeleteSheet("Sheet1")

	return f.SaveAs(filename)
}

// RunScraper 运行爬虫
func RunScraper() {
	config := DefaultConfig()

	// 合并策略ETF和基准ETF
	allETFs := append([]string{}, config.ETFCodes...)
	allETFs = append(allETFs, config.BenchmarkCodes...)

	isDaily := config.RebalanceMode == "daily"
	dataType := "周线"
	if isDaily {
		dataType = "日线"
	}

	fmt.Printf("开始爬取ETF%s数据...\n", dataType)
	fmt.Printf("策略ETF: %v\n", config.ETFCodes)
	fmt.Printf("基准ETF: %v\n", config.BenchmarkCodes)
	fmt.Printf("时间范围: %s 至 %s\n", config.StartDate, config.EndDate)

	scraper := NewETFScraper()
	if err := scraper.ScrapeAllETFs(allETFs, config.StartDate, config.EndDate, config.DataFile, isDaily); err != nil {
		fmt.Printf("错误: %v\n", err)
		return
	}

	fmt.Println("\n爬取完成！")
}
